import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import pandas as pd

# 参数
a = 55 / (2 * np.pi)  # 螺线的参数
v_a = 100  # a点的线速度，单位：cm/s
dt = 0.01  # 时间步长
t_max = 10  # 最大时间
fixed_distance = 165  # a和b之间的固定欧式距离

# 初始化
theta_a = 16 * 2 * np.pi  # 初始极角，设置在第10圈的位置

# 计算b点的初始极角，使得a和b之间的欧式距离为fixed_distance
theta_b = theta_a + np.arccos(1 - (fixed_distance ** 2) / (2 * a ** 2 * theta_a ** 2))

# 创建图形
fig, ax = plt.subplots()
line_a, = ax.plot([], [], 'ro', lw=2)  # a点用红色
line_b, = ax.plot([], [], 'bo', lw=2)  # b点用蓝色
line_ab, = ax.plot([], [], 'g-', lw=1)  # ab线段用绿色
ax.set_xlim(-1000, 1000)
ax.set_ylim(-1000, 1000)
ax.set_aspect('equal')

# 绘制等距螺线
theta_values = np.linspace(0, 50 * np.pi, 5000)
x_values = a * theta_values * np.cos(theta_values)
y_values = a * theta_values * np.sin(theta_values)
ax.plot(x_values, y_values, 'k--', lw=1)

# 记录坐标和速度
a_coords = []
b_coords = []
a_speeds = []
b_speeds = []


def init():
    line_a.set_data([], [])
    line_b.set_data([], [])
    line_ab.set_data([], [])
    return line_a, line_b, line_ab,


def animate(i):
    global theta_a, theta_b
    if theta_a > 6*2*np.pi:  # 当a点未到达原点时继续运动
        theta_a -= v_a / a * dt  # 向原点运动
        theta_b = theta_a + np.arccos(1 - (fixed_distance ** 2) / (2 * a ** 2 * theta_a ** 2))  # 保持固定欧式距离
    x_a = a * theta_a * np.cos(theta_a)
    y_a = a * theta_a * np.sin(theta_a)
    x_b = a * theta_b * np.cos(theta_b)
    y_b = a * theta_b * np.sin(theta_b)
    line_a.set_data(x_a, y_a)
    line_b.set_data(x_b, y_b)
    line_ab.set_data([x_a, x_b], [y_a, y_b])

    # 记录坐标
    a_coords.append((x_a, y_a))
    b_coords.append((x_b, y_b))

    # 计算速度
    if i > 0:
        # 计算a点的瞬时速度
        v_a_last = v_a
        # 计算b点的瞬时速度
        alpha = np.arctan2(y_b - y_a, x_b - x_a) - theta_a
        beta = np.arctan2(y_b - y_a, x_b - x_a) - theta_b
        v_b_last = v_a * np.cos(alpha) / np.cos(beta)
        a_speeds.append(v_a_last)
        b_speeds.append(v_b_last)

    return line_a, line_b, line_ab,


ani = animation.FuncAnimation(fig, animate, init_func=init, frames=int(t_max / dt), interval=20, blit=True)
plt.show()

# 输出坐标和速度至Excel文件
df = pd.DataFrame(columns=['0s', '1s', '2s', '3s', '4s', '5s', '6s', '7s', '8s', '9s', '10s'])

# 每秒的坐标和速度
for i in range(int(t_max) + 1):
    index = i * int(1 / dt)
    if index < len(a_coords):
        df.loc['a点x(m)', f'{i}s'] = a_coords[index][0] / 100
        df.loc['a点y(m)', f'{i}s'] = a_coords[index][1] / 100
        df.loc['b点x(m)', f'{i}s'] = b_coords[index][0] / 100
        df.loc['b点y(m)', f'{i}s'] = b_coords[index][1] / 100
        if index > 0:
            df.loc['a点速度(m/s)', f'{i}s'] = a_speeds[index - 1] / 100
            df.loc['b点速度(m/s)', f'{i}s'] = b_speeds[index - 1] / 100

df.to_excel('result1.xlsx', index=True)